import os
import pickle
import re
import string

import pandas as pd
from sklearn.preprocessing import LabelEncoder
from base.constant import out_folder, mapping_values
from base.util.file_util import check_path


def get_mapping_label_encoder(filename):
    if os.path.exists(filename):
        with open(filename, 'rb') as file:
            les = pickle.load(file)
    else:
        les = {}
    return les


# 保存LabelEncoder映射到本地
def save_mapping_label_encoder(les, filename):
    with open(filename, 'wb') as file:
        pickle.dump(les, file)


def getTokens(inputs):
    split_word = string.punctuation.replace('-', '').replace('_', '').replace('.', '')
    inputs = inputs.replace("\\", " ")
    tokensBySlash = re.sub(f"[{split_word}]", " ", inputs)
    allTokens = tokensBySlash.split(" ")
    allTokens = " ".join(sorted(set(allTokens)))
    return allTokens


def generate_mapping(config_yaml):
    if not config_yaml['mapping']['is_exec_mapping']:
        print("未启动映射,(若要启动，请添加is_exec_mapping配置)")
        return
    label_data = pd.read_csv(f"{out_folder}{os.sep}{config_yaml['label']['label_out']}",
                             delimiter=',')
    objects_num = label_data.select_dtypes(include=['number']).columns
    label_data[objects_num].fillna(0)
    label_data = label_data.rename(columns=mapping_values)
    objects = label_data.select_dtypes(include=['object']).columns
    les = get_mapping_label_encoder(config_yaml['mapping']['label_encoder_path'])
    for i, col in enumerate(objects):
        print(f'对象{i + 1}/{len(objects)}')
        label_data[col] = label_data[col].fillna('blank')
        label_data[col] = label_data[col].apply(lambda x: getTokens(str(x)))
        if col in les:
            le = les[col]
        else:
            le = LabelEncoder()
            les[col] = le
        label_data[col] = le.fit_transform(label_data[col])
        save_mapping_label_encoder(les, config_yaml['mapping']['label_encoder_path'])
    exceptObjEnum = label_data.select_dtypes(exclude=['object', 'bool']).columns
    for i, col in enumerate(exceptObjEnum):
        print(f'非对象枚举{i + 1}/{len(exceptObjEnum)}')
        label_data[col] = label_data[col].fillna(0)
    label_data.to_csv(
        check_path(f"{out_folder}{os.sep}{config_yaml['mapping']['mapping_out']}"),
        index=False)
    print("生成映射完毕，记得启动 arkime capture")
